import os
from openai import OpenAI


def get_api_key():
    """
    获取API密钥。

    Returns:
        str: API密钥。

    Raises:
        ValueError: 如果API密钥未找到。
    """
    # 从环境变量中获取API密钥
    key = os.getenv("BAICHUAN_API_KEY")
    # 如果密钥未找到，则抛出异常
    if not key:
        raise ValueError("API密钥未找到。请确保已设置环境变量BAICHUAN_API_KEY。")
    return key


def initialize_client(api_key):
    """
    初始化OpenAI客户端。

    Args:
        api_key (str): API密钥。

    Returns:
        OpenAI: 初始化的OpenAI客户端。
    """
    # 使用API密钥初始化OpenAI客户端
    return OpenAI(
        api_key=api_key,
        base_url="https://api.baichuan-ai.com/v1/",
    )


def send_message(client, model, messages):
    """
    发送消息并获取流式响应。

    Args:
        client (OpenAI): OpenAI客户端。
        model (str): 使用的模型名称。
        messages (list): 对话消息列表。

    Returns:
        str: 助手的完整回复。
    """
    # 初始化响应内容
    response_content = ""
    # 发送请求并获取流式响应
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0.3,
        stream=True,
    )
    # 处理流式响应
    for chunk in response:
        chunk_message = chunk.choices[0].delta
        if chunk_message.content:
            response_content += chunk_message.content
            print(chunk_message.content, end='', flush=True)
    return response_content


def main():
    """
    主函数。
    """
    # 获取API密钥
    api_key = get_api_key()
    # 初始化OpenAI客户端
    client = initialize_client(api_key)
    # 设置模型名称
    model = "Baichuan4-Air"
    # 初始对话历史
    messages = [
        {"role": "system", "content": "你是一个专业的SQL助手，根据我的需求回答问题。"}
    ]

    while True:
        # 发送消息并获取响应
        response_content = send_message(client, model, messages)
        # 将响应添加到对话历史
        messages.append({"role": "assistant", "content": response_content})
        # 获取用户输入
        user_input = input("\n你: ")
        # 如果用户输入 "exit" 或 "quit"，退出循环
        if user_input.lower() in ["exit", "quit"]:
            break
        # 将用户输入添加到对话历史
        messages.append({"role": "user", "content": user_input})


if __name__ == "__main__":
    main()
